import logging

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

import pyro
import pyro.distributions as dist
from pyro import poutine
from pyro.infer import SVI, Trace_ELBO
from pyro.infer.util import torch_item
from pyro.nn import PyroModule
from pyro.optim import ClippedAdam
from pyro.util import torch_isnan

logger = logging.getLogger(__name__)


class FullyConnected(nn.Sequential):
    """
    Fully connected multi-layer network with ELU activations.
    """
    # 初始化方法：默认添加一些全连接层和激活函数
    # sizes: 一个列表，定义每一层的输入和输出节点数。例如，[64, 128, 256] 表示输入维度为 64，第一层输出为 128，第二层输出为 256。
    def __init__(self, sizes, final_activation=None):
        layers = []
        # 遍历 sizes 列表的相邻元素对。
        # 例如，sizes = [64, 128, 256]，则 zip(sizes, sizes[1:]) 生成的对是 (64, 128) 和 (128, 256)。
        for in_size, out_size in zip(sizes, sizes[1:]):
            # 为每对 in_size 和 out_size 添加一个全连接层（nn.Linear）
            layers.append(nn.Linear(in_size, out_size))
            # 紧接着全连接层，添加 ELU 激活函数模块。
            layers.append(nn.ELU())

        # 移除最后一个添加的激活函数。
        # 因为在循环中每一层后都加了 ELU，但最后一层不一定需要激活函数，视用户需求而定。
        layers.pop(-1)
        # 检查是否传入了 final_activation。
        # 如果用户指定了一个最终的激活函数，就将其添加到网络的末尾。
        if final_activation is not None:
            layers.append(final_activation)

        # 调用父类 nn.Sequential 的构造函数，并将 layers 列表中的所有层传递进去。
        # *layers 会将列表展开成参数，传入 nn.Sequential，从而初始化一个按顺序连接的模块。
        super().__init__(*layers)

    # 定义了一个方法，用于向当前网络中动态添加新的层。
    def append(self, layer):
        # 确保传入的 layer 是一个 PyTorch 模块（nn.Module），否则会抛出错误。
        assert isinstance(layer, nn.Module)
        # 使用 add_module 方法将新的层添加到当前模型中。
        # str(len(self)) 将当前网络的长度（已包含的层数）作为新模块的名称，确保模块名称唯一。
        self.add_module(str(len(self)), layer)


class DistributionNet(nn.Module):
    """
    Base class for distribution nets.
    """
    # get_class: 静态方法，根据字符串动态获取子类。
    # 遍历 DistributionNet 的所有子类。
    # 比较子类名称与目标名称是否匹配（忽略大小写）。
    @staticmethod
    def get_class(dtype):
        """
        Get a subclass by a prefix of its name, e.g.::

            assert DistributionNet.get_class("bernoulli") is BernoulliNet
        """
        for cls in DistributionNet.__subclasses__():
            if cls.__name__.lower() == dtype + "net":
                return cls
        raise ValueError("dtype not supported: {}".format(dtype))


# 主要功能:
# 利用 FullyConnected 类，构造一个输出单个 logits 值的网络。
# 这个 logits 值可用于定义一个 Bernoulli 随机变量的条件概率分布。
# 示例:
# BernoulliNet([3, 4]) 定义一个两层的全连接网络，输入维度为 3，隐藏层维度为 4（为什么是隐藏层，因为最后一层是1），最终输出一个 logits 值。
# z 是一个形状为 (3,) 的输入向量。
# logits 是网络的输出，表示 Bernoulli 分布的对数几率。
# net.make_dist(logits).sample() 从生成的 Bernoulli 分布中采样。
class BernoulliNet(DistributionNet):
    """
    :class:`FullyConnected` network outputting a single ``logits`` value.

    This is used to represent a conditional probability distribution of a
    single Bernoulli random variable conditioned on a ``sizes[0]``-sized real
    value, for example::

        net = BernoulliNet([3, 4])
        z = torch.randn(3)
        logits, = net(z)
        t = net.make_dist(logits).sample()
    """
    def __init__(self, sizes):
        # 检查 sizes 列表的长度是否大于等于 1。
        # 确保网络至少有输入层。
        assert len(sizes) >= 1
        # 调用父类 DistributionNet 的构造函数。
        # 用于初始化继承自父类的属性或方法。
        super().__init__()
        # 使用 FullyConnected 构造一个多层全连接网络。
        # sizes + [1]: 将 [1] 追加到 sizes 的末尾，以确保网络的最终输出维度为 1（因为 Bernoulli 分布的 logits 是标量）。
        # 例如，sizes = [3, 4] 时，FullyConnected([3, 4, 1]) 表示输入 3，隐藏层输出 4，最终输出 1。
        self.fc = FullyConnected(sizes + [1])

    def forward(self, x):
        # logits = self.fc(x)
        # 将输入 x 通过全连接网络 self.fc。
        # 输出是一个形状为 (batch_size, 1) 的张量。
        # .squeeze(-1)
        # 移除最后一个维度（值为 1 的维度）。
        # 例如，如果 logits 的形状为 (batch_size, 1)，则 squeeze(-1) 后形状变为 (batch_size,)。
        # .clamp(min=-10, max=10)
        # 对 logits 值进行裁剪，使其范围限制在 [-10, 10] 之间。
        # 防止数值过大或过小，可能引起数值不稳定或梯度爆炸。
        logits = self.fc(x).squeeze(-1).clamp(min=-10, max=10)

        # 返回 logits，用逗号将其包装为元组形式。
        # 在 PyTorch 中，返回元组是为了兼容多输出的需求。
        return logits,

    # 定义一个方法，用于根据 logits 构造一个 Bernoulli 分布。
    @staticmethod
    def make_dist(logits):
        # 使用 torch.distributions.Bernoulli 创建一个 Bernoulli 分布对象。
        # logits: 定义 Bernoulli 分布的对数几率（logits）。
        # logits 是概率 p 的对数几率
        # 返回的分布对象支持采样、计算概率等操作。
        return dist.Bernoulli(logits=logits)


# 主要功能:
# 基于 FullyConnected 网络生成受约束的 rate 参数。
# rate 用于定义指数分布（Exponential distribution）。
# 示例用法:
# 创建网络 net = ExponentialNet([3, 4])，输入维度为 3，隐藏层维度为 4，最终输出一个 rate 值。
# 输入向量 x 的维度为 (3,)。
# 网络输出 rate 值，通过 net.make_dist(rate) 构造指数分布，从中采样出随机变量 y。
# 总结
# ExponentialNet 用于构造基于输入特征的条件指数分布：
# 前向传播生成受约束的 rate 参数，保证其正值范围合理。
# 提供 make_dist 方法，用于根据 rate 参数构造 Exponential 分布对象。
# 典型用途包括模拟等待时间分布、事件发生的时间间隔等场景。
class ExponentialNet(DistributionNet):
    """
    :class:`FullyConnected` network outputting a constrained ``rate``.

    This is used to represent a conditional probability distribution of a
    single Normal random variable conditioned on a ``sizes[0]``-size real
    value, for example::

        net = ExponentialNet([3, 4])
        x = torch.randn(3)
        rate, = net(x)
        y = net.make_dist(rate).sample()
    """
    def __init__(self, sizes):
        assert len(sizes) >= 1
        super().__init__()
        self.fc = FullyConnected(sizes + [1])

    def forward(self, x):
        # self.fc(x)
        #   将输入 x 通过全连接网络 self.fc。
        #   输出是形状为 (batch_size, 1) 的张量。
        # .squeeze(-1)
        #   移除最后一个维度，使输出从 (batch_size, 1) 变为 (batch_size,)。
        # nn.functional.softplus()
        #   对输出值应用 Softplus 函数。
        # Softplus 是一种平滑的 ReLU 函数，定义为：
        # Softplus 可以确保输出值为正，因此适合作为正数约束（如指数分布的 scale）。
        # .clamp(min=1e-3, max=1e6)
        #   将 Softplus 的输出值裁剪到范围 [1e-3, 1e6]。
        #   目的:
        #   min=1e-3: 防止数值过小，避免数值不稳定。
        #   max=1e6: 限制过大的数值，防止梯度爆炸。
        # .reciprocal()
        #   计算 scale 的倒数，得到 rate。
        # 在指数分布中，rate 是分布的参数，定义为：

        # return rate,
        #   返回 rate 值，使用逗号将其包装为元组形式。
        #   在 PyTorch 中，返回元组可以兼容多输出的需求。
        scale = nn.functional.softplus(self.fc(x).squeeze(-1)).clamp(min=1e-3, max=1e6)
        rate = scale.reciprocal()
        return rate,

    @staticmethod
    def make_dist(rate):
        # 使用 torch.distributions.Exponential 创建一个指数分布对象。
        # 参数：
        # rate: 指数分布的速率参数。
        return dist.Exponential(rate)


# 该网络用于生成 loc（位置参数）和 scale（尺度参数）来构造拉普拉斯分布（Laplace distribution）
# 通过 FullyConnected 网络输出 loc 和 scale 参数。
# loc 和 scale 约束在一定范围内，避免数值不稳定。
# 示例用法：
# 创建网络 net = LaplaceNet([3, 4])，输入维度为 3，隐藏层维度为 4，最终输出 loc 和 scale 两个参数。
# 输入 x 形状为 (3,)，网络计算出 loc, scale。
# make_dist(loc, scale) 生成拉普拉斯分布，并从中采样 y。
# 总结
# LaplaceNet 通过 FullyConnected 网络生成 loc（位置参数）和 scale（尺度参数）。
# 通过 softplus 约束 scale 为正，确保数值稳定性。
# 提供 make_dist(loc, scale) 方法，根据 loc 和 scale 创建 torch.distributions.Laplace 对象。
# 可用于建模具有尖峰和长尾分布的随机变量（如噪声建模、稳健回归等）。
class LaplaceNet(DistributionNet):
    """
    :class:`FullyConnected` network outputting a constrained ``loc,scale``
    pair.

    This is used to represent a conditional probability distribution of a
    single Laplace random variable conditioned on a ``sizes[0]``-size real
    value, for example::

        net = LaplaceNet([3, 4])
        x = torch.randn(3)
        loc, scale = net(x)
        y = net.make_dist(loc, scale).sample()
    """
    def __init__(self, sizes):
        assert len(sizes) >= 1
        super().__init__()
        # 创建一个 FullyConnected 网络：
        # sizes + [2] 在 sizes 列表末尾追加 [2]，表示最终输出维度为 2，分别对应 loc 和 scale 参数。
        self.fc = FullyConnected(sizes + [2])

    def forward(self, x):
        # 通过 FullyConnected 网络 self.fc 计算 loc_scale，其形状为 (batch_size, 2)。
        loc_scale = self.fc(x)
        # 提取 loc_scale 的第 0 维分量（即 loc）。
        # 使用 .clamp(min=-1e6, max=1e6) 限制 loc 的取值范围：
        # min=-1e6: 防止 loc 过小，避免溢出。
        # max=1e6: 防止 loc 过大，避免梯度爆炸。
        loc = loc_scale[..., 0].clamp(min=-1e6, max=1e6)
        # 提取 loc_scale 的第 1 维分量（即 scale）。
        # 使用 softplus 函数确保 scale 为正
        # .clamp(min=1e-3, max=1e6) 进一步限制 scale 范围：
        # min=1e-3: 避免 scale 过小，防止数值不稳定。
        # max=1e6: 避免 scale 过大，防止梯度爆炸。
        scale = nn.functional.softplus(loc_scale[..., 1]).clamp(min=1e-3, max=1e6)
        return loc, scale

    @staticmethod
    def make_dist(loc, scale):
        # 使用 torch.distributions.Laplace 生成一个拉普拉斯分布对象：
        # loc 代表分布的中心位置。
        # scale 代表分布的尺度（决定了数据的离散程度）。
        return dist.Laplace(loc, scale)


# 该网络用于生成 loc（均值）和 scale（标准差），以构造正态分布（Normal distribution）。
# 通过 FullyConnected 网络输出 loc 和 scale 参数。
# loc 代表正态分布的均值，scale 代表标准差（必须为正）。
# 示例用法：
# 创建 NormalNet 网络：
#   net = NormalNet([3, 4])
# 输入一个 torch.randn(3) 张量，计算 loc, scale。
# 使用 make_dist(loc, scale) 生成正态分布，并采样 y。
class NormalNet(DistributionNet):
    """
    :class:`FullyConnected` network outputting a constrained ``loc,scale``
    pair.

    This is used to represent a conditional probability distribution of a
    single Normal random variable conditioned on a ``sizes[0]``-size real
    value, for example::

        net = NormalNet([3, 4])
        x = torch.randn(3)
        loc, scale = net(x)
        y = net.make_dist(loc, scale).sample()
    """
    def __init__(self, sizes):
        assert len(sizes) >= 1
        super().__init__()
        self.fc = FullyConnected(sizes + [2])

    def forward(self, x):
        # 通过 FullyConnected 网络 self.fc 计算 loc_scale，其形状为 (batch_size, 2)。
        # 其中：
        # loc_scale[..., 0] 对应 loc（均值）。
        # loc_scale[..., 1] 对应 scale（标准差）。
        loc_scale = self.fc(x)
        # 提取 loc_scale 的第 0 维分量（即 loc）。
        # 使用 .clamp(min=-1e6, max=1e6) 限制 loc 的取值范围：
        # min=-1e6: 防止 loc 过小，避免数值溢出。
        # max=1e6: 防止 loc 过大，避免梯度爆炸。
        loc = loc_scale[..., 0].clamp(min=-1e6, max=1e6)
        # 提取 loc_scale 的第 1 维分量（即 scale）。
        # 使用 softplus 函数确保 scale 为正
        # .clamp(min=1e-3, max=1e6) 进一步限制 scale：
        # min=1e-3: 避免 scale 过小，防止数值不稳定。
        # max=1e6: 避免 scale 过大，防止梯度爆炸。
        scale = nn.functional.softplus(loc_scale[..., 1]).clamp(min=1e-3, max=1e6)
        return loc, scale

    @staticmethod
    def make_dist(loc, scale):
        return dist.Normal(loc, scale)


# 该网络用于生成 df（自由度）、loc（均值）、scale（标准差），以构造 Student-T 分布。
# 通过 FullyConnected 网络输出 loc 和 scale 参数，并额外包含 df（自由度）。
# df > 1，因为 df=1 代表柯西分布，df→∞ 时趋近于正态分布。
# 示例用法：
# 创建 StudentTNet 网络：
# net = StudentTNet([3, 4])
# 输入 x = torch.randn(3)，计算 df, loc, scale。
# 使用 make_dist(df, loc, scale) 生成 Student-T 分布，并采样 y。
class StudentTNet(DistributionNet):
    """
    :class:`FullyConnected` network outputting a constrained ``df,loc,scale``
    triple, with shared ``df > 1``.

    This is used to represent a conditional probability distribution of a
    single Student's t random variable conditioned on a ``sizes[0]``-size real
    value, for example::

        net = StudentTNet([3, 4])
        x = torch.randn(3)
        df, loc, scale = net(x)
        y = net.make_dist(df, loc, scale).sample()
    """
    def __init__(self, sizes):
        assert len(sizes) >= 1
        super().__init__()
        self.fc = FullyConnected(sizes + [2])
        # df_unconstrained 是一个 可学习参数（自由度 df 的未约束版本）。
        # 直接初始化为 0.0，之后通过 softplus 转换为正值，确保 df > 1。
        self.df_unconstrained = nn.Parameter(torch.tensor(0.0))

    def forward(self, x):
        # 通过 FullyConnected 网络 self.fc 计算 loc_scale，其形状为 (batch_size, 2)。
        # 其中：
        # loc_scale[..., 0] 对应 loc（均值）。
        # loc_scale[..., 1] 对应 scale（标准差）。
        loc_scale = self.fc(x)
        loc = loc_scale[..., 0].clamp(min=-1e6, max=1e6)
        scale = nn.functional.softplus(loc_scale[..., 1]).clamp(min=1e-3, max=1e6)
        # 计算 df（自由度）：
        # 先对 self.df_unconstrained 应用 softplus，确保 df 为正： df = \log(1 + e^{\text{df_unconstrained}}) + 1
        # +1 确保 df > 1，避免过小的自由度导致数值不稳定。
        # .expand_as(loc) 让 df 具有与 loc 相同的形状（batch-wise 共享相同的 df）。
        df = nn.functional.softplus(self.df_unconstrained).add(1).expand_as(loc)
        return df, loc, scale

    @staticmethod
    def make_dist(df, loc, scale):
        return dist.StudentT(df, loc, scale)


# 该网络用于输出 均值（loc） 和 标准差（scale），从而构造一个对角协方差矩阵的多维正态分布。
# DiagNormalNet 通过 FullyConnected 网络计算 loc（均值）和 scale（标准差）。
# 生成一个 sizes[-1] 维的 对角正态分布（即，每个维度的变量独立）。
# 示例用法：
# net = DiagNormalNet([3, 4, 5])  # 创建一个 [3, 4, 5] 结构的网络
# z = torch.randn(3)  # 输入 3 维随机数
# loc, scale = net(z)  # 计算 loc 和 scale
# x = dist.Normal(loc, scale).sample()  # 采样 x
# 适用于：
# 潜变量（latent z）建模。
# 标准化（prewhitened）特征 x。
class DiagNormalNet(nn.Module):
    """
    :class:`FullyConnected` network outputting a constrained ``loc,scale``
    pair.

    This is used to represent a conditional probability distribution of a
    ``sizes[-1]``-sized diagonal Normal random variable conditioned on a
    ``sizes[0]``-size real value, for example::

        net = DiagNormalNet([3, 4, 5])
        z = torch.randn(3)
        loc, scale = net(z)
        x = dist.Normal(loc, scale).sample()

    This is intended for the latent ``z`` distribution and the prewhitened
    ``x`` features, and conservatively clips ``loc`` and ``scale`` values.
    """
    def __init__(self, sizes):
        # 确保 sizes 至少包含 2 个元素，即：
        # sizes[0] 代表输入维度。
        # sizes[-1] 代表最终的分布维度（正态分布的变量个数）。
        assert len(sizes) >= 2
        # 记录 dim（最终正态分布的维度），即 sizes 最后一维的大小。
        self.dim = sizes[-1]
        super().__init__()
        # 创建 全连接神经网络（FullyConnected）：
        # sizes[:-1] 去掉 sizes 的最后一维，作为隐藏层的配置。
        # self.dim * 2 作为最后的输出维度：
        # dim 维度的 loc（均值）。
        # dim 维度的 scale（标准差）。
        self.fc = FullyConnected(sizes[:-1] + [self.dim * 2])

    def forward(self, x):
        # 通过 FullyConnected 网络 self.fc 计算 loc_scale，形状为 (batch_size, dim * 2)。
        # 其中：
        # loc_scale[..., :dim] 对应 loc（均值）。
        # loc_scale[..., dim:] 对应 scale（标准差）。
        loc_scale = self.fc(x)
        # 提取 loc_scale 的前 dim 维，作为 loc。
        # 使用 .clamp(min=-1e2, max=1e2) 限制 loc 的范围：
        # min=-1e2: 防止 loc 过小，避免数值不稳定。
        # max=1e2: 防止 loc 过大，避免梯度爆炸。
        loc = loc_scale[..., :self.dim].clamp(min=-1e2, max=1e2)
        # 提取 loc_scale 的后 dim 维，作为 scale。
        # 使用 softplus 确保 scale 为正
        # .add(1e-3) 确保 scale 最小值不为 0，防止 scale 过小导致数值不稳定。
        # .clamp(max=1e2) 限制 scale 的最大值，避免梯度爆炸。
        scale = nn.functional.softplus(loc_scale[..., self.dim:]).add(1e-3).clamp(max=1e2)
        return loc, scale


# 定义 DiagBernoulliNet 类，继承 torch.nn.Module，用于建模 Bernoulli 分布（即二元随机变量的概率分布）。
# logits 是 Bernoulli 分布 的 对数几率（logits），可以转换为概率
# 通过 FullyConnected 网络计算 logits。
# logits 可以用于构造 Bernoulli 分布。
# 示例用法：
# net = DiagBernoulliNet([3, 4, 5])
# z = torch.randn(3)  # 输入 3 维随机数
# logits, = net(z)  # 计算 logits
# t = net.make_dist(logits).sample()  # 采样一个 Bernoulli 变量

# 总结
# DiagBernoulliNet 主要用于建模 dim 维的 Bernoulli 分布：
# logits（对数几率）是网络输出。
# 通过 sigmoid(logits) 计算 概率，用于 Bernoulli 分布采样 0/1 结果。
# 特点：
# logits 通过 FullyConnected 网络计算。
# clamp(0, 11) 限制 logits，提高数值稳定性。
# make_dist(logits) 用于构造 torch.distributions.Bernoulli 分布。
# 应用场景：
# 离散二元变量建模（如二分类任务）。
# 概率图模型中的离散采样（如 VAE 变分自编码器中的离散变量）。
class DiagBernoulliNet(nn.Module):
    """
    :class:`FullyConnected` network outputting a single ``logits`` value.

    This is used to represent a conditional probability distribution of a
    single Bernoulli random variable conditioned on a ``sizes[0]``-sized real
    value, for example::

        net = DiagBernoulliNet([3, 4, 5])
        z = torch.randn(3)
        logits, = net(z)
        t = net.make_dist(logits).sample()
    """
    def __init__(self, sizes):
        # 确保 sizes 至少包含 2 个元素，即：
        # sizes[0] 代表输入维度。
        # sizes[-1] 代表输出 logits 的维度。
        assert len(sizes) >= 2
        # 记录 dim（最终 logits 的维度），即 sizes 最后一维的大小。
        self.dim = sizes[-1]
        super().__init__()
        # 创建 全连接神经网络（FullyConnected）：
        # sizes[:-1] 去掉 sizes 的最后一维，作为隐藏层的配置。
        # self.dim 作为最终输出维度：
        # dim 维度的 logits，用于构造 Bernoulli 分布。
        # 例如：
        # net = DiagBernoulliNet([3, 4, 5])
        # 这里 FullyConnected([3, 4, 5])：
        # 输入维度 3
        # 隐藏层维度 4
        # 输出维度 5（对应 5 个 Bernoulli 变量的 logits）。
        self.fc = FullyConnected(sizes[:-1] + [self.dim])

    def forward(self, x):
        # 通过 FullyConnected 网络 self.fc 计算 logits，形状为 (batch_size, dim)。
        # logits 代表 Bernoulli 分布的对数几率（logit）：
        # 可以通过 sigmoid(logits) 计算得到概率
        # .squeeze(-1)
        # 如果 logits 形状是 (batch_size, dim, 1)，那么 .squeeze(-1) 会去掉 1 这个维度，确保形状变成 (batch_size, dim)。
        # .clamp(min=0, max=11)
        # 限制 logits 的取值范围：
        # min=0：防止 logits 过小，使得 sigmoid(logits) 计算的概率不会过低。
        # max=11：防止 logits 过大，避免梯度爆炸。
        logits = self.fc(x).squeeze(-1).clamp(min=0, max=11)
        return logits

    @staticmethod
    def make_dist(logits):
        return dist.Bernoulli(logits=logits)


# 定义 PreWhitener 类，继承 torch.nn.Module。
# 功能：执行 数据白化（whitening） 预处理，即 去均值 并 标准化方差，使数据变得更稳定。
class PreWhitener(nn.Module):
    """
    Data pre-whitener.
    """
    def __init__(self, data):
        super().__init__()
        with torch.no_grad():
            # 计算 data 按列（特征维度） 的 均值，得到形状 (D,) 的向量。
            # loc 代表数据的 中心位置，用于 去均值 处理
            loc = data.mean(0)
            # 计算 data 按列（特征维度） 的 标准差，得到形状 (D,) 的向量
            # scale 用于标准化数据（除以标准差）
            scale = data.std(0)
            # 处理可能出现的 标准差为 0 的情况：
            # 如果 scale[j] <= 0，说明某个特征的所有样本值都相同（无变化）。
            # 这样会导致除以 scale[j] 时报错，因此将其设为 1.0
            scale[~(scale > 0)] = 1.
            # 使用 register_buffer 注册 loc（均值向量）为 buffer：
            # 不会被优化器更新，但会随着模型保存和加载。
            # loc 用于去均值处理。
            self.register_buffer("loc", loc)
            # 计算 inv_scale = 1 / scale（元素级倒数）。
            # 避免除法计算，提高数值稳定性
            # 将 inv_scale 也注册为 buffer，保证它随着模型保存和加载。
            self.register_buffer("inv_scale", scale.reciprocal())

    def forward(self, data):
        # (data - self.loc)
        # 去均值（mean subtraction），使数据中心化
        # self.loc 形状 (D,)，自动 广播（broadcasting）到 (N, D)。
        # * self.inv_scale
        # 标准化（scale normalization），使数据 标准差归一
        # self.inv_scale 形状 (D,)，自动 广播（broadcasting）到 (N, D)。
        return (data - self.loc) * self.inv_scale


# 这个 Guide 类是一个用于 因果效应估计（causal effect estimation） 的推断模型。它使用 变分推断（variational inference） 方法来学习潜在混杂变量 z 以及处理二值处理变量 t 和结果变量 y 的条件分布。
# Guide 负责定义推断模型，用于估计因果效应。
# z、zt、zy 通过 x 计算并采样。
# t 由 z 和 zt 计算并采样。
# y 由 t、z 和 zy 计算并采样。
# Guide 通过神经网络学习这些分布的参数。
class Guide(PyroModule):
    """
    Inference model for causal effect estimation with latent confounder ``z``
    and binary treatment ``t``::
        z ~ p(z|x)  # latent confounder, an embedding
        zt ~ p(zt|x)
        zy ~ p(zy|x)
        t ~ p(t|z,zt) # treatment
        y ~ p(y|t,z,zy)    # outcome

    Each of these distributions is defined by a neural network.  The ``y`` and
    ``z`` distributions are defined by disjoint pairs of neural networks
    defining ``p(-|t=0,...)`` and ``p(-|t=1,...)``; this allows highly
    imbalanced treatment.

    :param dict config: A dict specifying ``feature_dim``, ``latent_dim``,
        ``hidden_dim``, ``num_layers``, and ``outcome_dist``.
    """
    def __init__(self, config):
        # 从 config 读取：
        # latent_dim：潜在混杂变量 z 的维度
        # latent_dim_t：用于处理 t 的潜在变量 zt 维度
        # latent_dim_y：用于处理 y 的潜在变量 zy 维度
        self.latent_dim = config["latent_dim"]
        self.latent_dim_t = config["latent_dim_t"]
        self.latent_dim_y = config["latent_dim_y"]

        # OutcomeNet 是 y 的分布网络。
        # config["outcome_dist"] 指定 y 的分布类型，DistributionNet.get_class(...) 负责返回相应的分布网络类。
        OutcomeNet = DistributionNet.get_class(config["outcome_dist"])

        super().__init__()
        # self.t_nn = BernoulliNet([config["feature_dim"]])

        # t_nn 用于建模 p(t|z, zt)（t 依赖于 z 和 zt）。
        # BernoulliNet(...) 负责输出 t 的伯努利分布参数（logits）。
        self.t_nn = BernoulliNet([config["latent_dim"]+config["latent_dim_t"]])

        # The y and z networks both follow an architecture where the first few
        # layers are shared for t in {0,1}, but the final layer is split
        # between the two t values.
        # y_nn 是一个全连接神经网络：
        #   输入 z 和 zy
        #   经过 num_layers-1 层隐藏层
        #   使用 ELU 作为激活函数
        self.y_nn = FullyConnected([config["latent_dim"] + config["latent_dim_y"]] +
                                   [config["hidden_dim"]] * (config["num_layers"] - 1),
                                   final_activation=nn.ELU())
        # y0_nn 和 y1_nn 分别表示 p(y|t=0, z, zy) 和 p(y|t=1, z, zy) 的参数化模型。
        self.y0_nn = OutcomeNet([config["hidden_dim"]])
        self.y1_nn = OutcomeNet([config["hidden_dim"]])

        # z_nn：输入 x，输出潜在变量 z 的隐藏表示。
        self.z_nn = FullyConnected([config["feature_dim"]] +
                                   [config["hidden_dim"]] * (config["num_layers"] - 1),
                                   final_activation=nn.ELU())
        
        # z_out_nn：将 z_nn 的输出映射到高斯分布 p(z|x) 的均值和标准差。
        self.z_out_nn = DiagNormalNet([config["hidden_dim"], config["latent_dim"]])
        
        # zt_nn 处理 p(zt|x)
        self.zt_nn = FullyConnected([config["feature_dim"]] +
                                   [config["hidden_dim"]] * (config["num_layers"] - 1),
                                   final_activation=nn.ELU())
        # zt_out_nn 计算 zt 服从的高斯分布参数
        self.zt_out_nn = DiagNormalNet([config["hidden_dim"], config["latent_dim_t"]])
        
        # zy_nn 处理 p(zy|x)
        self.zy_nn = FullyConnected([config["feature_dim"]] +
                                   [config["hidden_dim"]] * (config["num_layers"] - 1),
                                   final_activation=nn.ELU())
        # zy_out_nn 计算 zy 服从的高斯分布参数
        self.zy_out_nn = DiagNormalNet([config["hidden_dim"], config["latent_dim_y"]])

    def forward(self, x, t=None, y=None, size=None):
        # x：特征输入
        # t：（可选）观测到的 t
        # y：（可选）观测到的 y
        # size：子采样大小，默认取 x.size(0)
        if size is None:
            size = x.size(0)
        # pyro.plate("data", size, subsample=x) 用于变分推断中的子采样（mini-batch training）。
        with pyro.plate("data", size, subsample=x):
            # The t and y sites are needed for prediction, and participate in
            # the auxiliary CEVAE loss. We mark them auxiliary to indicate they
            # do not correspond to latent variables during training.
            # 采样潜在变量 z、zt 和 zy。
            z=pyro.sample("z", self.z_dist(x))
            zt=pyro.sample("zt", self.zt_dist(x))
            zy=pyro.sample("zy", self.zy_dist(x))

            # t 由 p(t|z,zt) 采样，如果 t 有观测值，则用于监督训练。
            t = pyro.sample("t", self.t_dist(z,zt), obs=t, infer={"is_auxiliary": True})

            # y 由 p(y|t,z,zy) 采样。
            y = pyro.sample("y", self.y_dist(t,z,zy), obs=y, infer={"is_auxiliary": True})
            # The z site participates only in the usual ELBO loss.

    # 用于计算 z、zt、zy 的均值。
    def z_mean(self, x, t=None):
        with pyro.plate("data", x.size(0)):
            z = pyro.sample("z", self.z_dist(x))
            zt = pyro.sample("zt", self.zt_dist(x))
            zy = pyro.sample("zy", self.zy_dist(x))
        return z,zt,zy

    # t 由 z 和 zt 计算得到的 logits 经过 Bernoulli 采样。
    def t_dist(self, z, zt):
        input_concat = torch.cat((z,zt),-1)
        logits, = self.t_nn(input_concat)
        return dist.Bernoulli(logits=logits)

    # 计算 y 的分布，t=0 和 t=1 分别对应 y0_nn 和 y1_nn。
    def y_dist(self, t, z, zy):
        # The first n-1 layers are identical for all t values.
        x = torch.cat((z,zy),-1)
        hidden = self.y_nn(x)
        # In the final layer params are not shared among t values.
        params0 = self.y0_nn(hidden)
        params1 = self.y1_nn(hidden)
        t = t.bool()
        params = [torch.where(t, p1, p0) for p0, p1 in zip(params0, params1)]
        return self.y0_nn.make_dist(*params)

    # z_dist、zt_dist、zy_dist 用 FullyConnected 计算高斯分布参数，返回 dist.Normal。
    def z_dist(self, x):
        # hidden = self.z_nn(x)
        hidden = self.z_nn(x.float())
        params = self.z_out_nn(hidden)
        return dist.Normal(*params).to_event(1)

    def zt_dist(self, x):
        hidden = self.zt_nn(x.float())
        params = self.zt_out_nn(hidden)
        return dist.Normal(*params).to_event(1)
    
    def zy_dist(self, x):
        hidden = self.zy_nn(x.float())
        params = self.zy_out_nn(hidden)
        return dist.Normal(*params).to_event(1)


#这段代码定义了一个新的损失函数 TraceCausalEffect_ELBO，它继承自 Trace_ELBO，用于训练 TEDVAE（一种用于因果效应估计的变分自编码器）。
# 其目标是在标准的 ELBO（证据下界）损失的基础上，额外考虑 q(t|z,zt) 和 q(y|t,z,zy)，以增强对潜在混杂因素的建模能力。
# 该损失函数用于 TEDVAE 训练。
# CEVAE（Causal Effect Variational Autoencoder）是一种因果推断模型。
# ELBO 是标准变分推断目标，q(t|z,zt) 和 q(y|t,z,zy) 是额外的对数概率项。
# 训练时，我们希望 最大化 -loss
# q(t|z,zt) 表示给定潜在变量 z 和 zt，估计 t（处理变量）的概率。
# q(y|t,z,zy) 表示给定 t, z, zy，估计 y（结果变量）的概率。
class TraceCausalEffect_ELBO(Trace_ELBO):
    """
    Loss function for training a :class:`TEDVAE`.
    From [1], the CEVAE objective (to maximize) is::

        -loss = ELBO + log q(t|z,zt) + log q(y|t,z,zy)
    """
    # 该方法计算单个样本粒子的损失，即 ELBO 和额外的对数概率项。
    def _differentiable_loss_particle(self, model_trace, guide_trace):
        # Construct -ELBO part.
        # guide_trace 是 Pyro 生成的计算图（即神经网络的前向传播过程），其中包含了所有随机变量的采样信息。
        # 遍历 guide_trace.nodes，筛选出 is_observed=True 的变量，即 观测到的变量（t 和 y），存入 blocked_names。
        # 这些观测变量 不会 作为隐变量的一部分参与 ELBO 计算，而是单独加入 log q 项。
        blocked_names = [name for name, site in guide_trace.nodes.items()
                         if site["type"] == "sample" and site["is_observed"]]

        # 复制 guide_trace，并 删除 blocked_names 中的节点（即 t 和 y），以便计算标准 ELBO。
        # 这样，在 super()._differentiable_loss_particle() 计算 ELBO 时，它不会将 t 和 y 视为潜在变量。
        blocked_guide_trace = guide_trace.copy()
        for name in blocked_names:
            del blocked_guide_trace.nodes[name]
        # 这一步计算 标准 ELBO（去除了 t 和 y ）。
        # loss 是最终优化的目标，surrogate_loss 是用于梯度计算的松弛损失。
        loss, surrogate_loss = super()._differentiable_loss_particle(
            model_trace, blocked_guide_trace)

        # Add log q terms.
        # 遍历 blocked_names（即 t 和 y），取出 log_prob_sum（即 q(t|z,zt) 和 q(y|t,z,zy) 的对数概率）。
        # 乘以 100 作为权重，使这两个额外项的影响更大。
        # 为何取 -100 * log_q？
        # log_q 是 q(t|z,zt) 和 q(y|t,z,zy) 的对数概率，我们希望 最大化 它们，因此在 损失函数中取负。
        for name in blocked_names:
            log_q = guide_trace.nodes[name]["log_prob_sum"]
            loss = loss - 100* torch_item(log_q)
            surrogate_loss = surrogate_loss - 100* log_q

        return loss, surrogate_loss

    @torch.no_grad()
    def loss(self, model, guide, *args, **kwargs):
        # 计算 全批次的损失（包括 ELBO 和 log_q 项）。
        # torch_item() 将张量转换为 Python 标量，方便返回。
        return torch_item(self.differentiable_loss(model, guide, *args, **kwargs))


# 实现了一个名为TEDVAE的模型，用于估计因果效应，特别是处理个体治疗效应(ITE)。
class Model(PyroModule):
    """
    Generative model for a causal model with latent confounder ``z`` and binary
    treatment ``t``::

        z ~ p(z)      # latent confounder
        zt ~ p(zt)
        zy ~ p（zy)
        x ~ p(x|z,zt,zy)
        t ~ p(t|z,zt)
        y ~ p(y|t,z,zy)

    Each of these distributions is defined by a neural network.  The ``y``
    distribution is defined by a disjoint pair of neural networks defining
    ``p(y|t=0,z,zy)`` and ``p(y|t=1,z,zy)``; this allows highly imbalanced treatment.

    :param dict config: A dict specifying ``feature_dim``, ``latent_dim``,
        ``hidden_dim``, ``num_layers``, and ``outcome_dist``.
    """
    def __init__(self, config):
        # 从配置中获取各种维度参数：潜在变量维度、治疗相关潜在维度、结果相关潜在维度、二进制特征维度和连续特征维度
        self.latent_dim = config["latent_dim"]
        self.latent_dim_t = config["latent_dim_t"]
        self.latent_dim_y = config["latent_dim_y"]
        self.binfeats = config["binary_dim"]
        self.contfeats = config["continuous_dim"]

        super().__init__()

        # 用于建模 p(x_cont|z,zt,zy) 和 p(x_bin|z,zt,zy) 的网络。
        # 初始化连续特征x的神经网络，输出对角正态分布参数
        # 网络结构：输入层(潜在变量拼接)→多个隐藏层→输出层(连续特征维度)
        self.x_nn = DiagNormalNet([config["latent_dim"]+config["latent_dim_t"]+config["latent_dim_y"]] +
                                  [config["hidden_dim"]] * config["num_layers"] +
                                  [len(config["continuous_dim"])])
        # 初始化二进制特征x的神经网络，输出伯努利分布参数
        # 网络结构与上面类似，但输出维度是二进制特征数量
        self.x2_nn = DiagBernoulliNet([config["latent_dim"]+config["latent_dim_t"]+config["latent_dim_y"]] +
                                  [config["hidden_dim"]] * config["num_layers"] +
                                  [len(config["binary_dim"])])
        # 获取结果分布的神经网络类
        OutcomeNet = DistributionNet.get_class(config["outcome_dist"])
        # The y network is split between the two t values.
        # 分别是 p(y|t=0,z,zy) 和 p(y|t=1,z,zy) 的两个网络，分别建模两个处理变量的输出。
        # 初始化两个结果网络(y0和y1)，分别对应t=0和t=1的情况
        # 网络输入是潜在变量z和zy的拼接
        self.y0_nn = OutcomeNet([config["latent_dim"]+config["latent_dim_y"]] +
                                [config["hidden_dim"]] * config["num_layers"])
        self.y1_nn = OutcomeNet([config["latent_dim"]+config["latent_dim_y"]] +
                                [config["hidden_dim"]] * config["num_layers"])
        # 建模 p(t|z,zt)。
        # 初始化治疗变量t的神经网络，输出伯努利分布参数
        # 输入是潜在变量z和zt的拼接
        self.t_nn = BernoulliNet([config["latent_dim"]+config["latent_dim_t"]])

    # x 为观测特征，t 为处理，y 为结果。size 是总数据数量（用于 ELBO 权重缩放）
    def forward(self, x, t=None, y=None, size=None):
        if size is None:
            size = x.size(0)
        # plate用于批处理和数据子采样
        with pyro.plate("data", size, subsample=x):
            # 从标准正态分布中采样潜在变量z、zt和zy
            z = pyro.sample("z", self.z_dist())
            zt = pyro.sample("zt", self.zt_dist())
            zy = pyro.sample("zy", self.zy_dist())

            # 分别生成二进制和连续特征，并将观测数据x作为条件
            # x = pyro.sample("x", self.x_dist(z, zt, zy), obs=x)
            x_binary = pyro.sample("x_bin", self.x_dist_binary(z, zt, zy), obs=x[:,self.binfeats])
            x_continuous = pyro.sample("x_cont", self.x_dist_continuous(z, zt, zy), obs=x[:,self.contfeats])
            # 将两部分特征拼接回完整x
            x = torch.cat((x_binary,x_continuous), -1)
            # x = pyro.sample("x", self.x_dist_binary(z, zt, zy), obs=x)
            # 生成治疗变量t和结果变量y，可以传入观测值
            t = pyro.sample("t", self.t_dist(z, zt), obs=t)
            y = pyro.sample("y", self.y_dist(t, z, zy), obs=y)

        return y

    # 计算y的期望值
    # 类似forward但不采样y，只计算均值
    # 用于计算因果推断中 E[y|do(t)]，在估计 ITE 时使用。
    def y_mean(self, x, t=None):
        with pyro.plate("data", x.size(0)):
            z = pyro.sample("z", self.z_dist())
            zt = pyro.sample("zt", self.zt_dist())
            zy = pyro.sample("zy", self.zy_dist())
            x_binary = pyro.sample("x_bin", self.x_dist_binary(z, zt, zy), obs=x[:,self.binfeats])
            x_continuous = pyro.sample("x_cont", self.x_dist_continuous(z, zt, zy), obs=x[:,self.contfeats])
            x = torch.cat((x_binary,x_continuous), -1)
            t = pyro.sample("t", self.t_dist(z, zt), obs=t)
        return self.y_dist(t, z, zy).mean

    # 定义潜在变量的先验分布(标准正态)
    def z_dist(self):
        return dist.Normal(0, 1).expand([self.latent_dim]).to_event(1)
    
    def zt_dist(self):
        return dist.Normal(0, 1).expand([self.latent_dim_t]).to_event(1)

    def zy_dist(self):
        return dist.Normal(0, 1).expand([self.latent_dim_y]).to_event(1)

    # 返回连续特征的条件分布
    def x_dist_continuous(self, z, zt, zy):
        z_concat = torch.cat((z,zt,zy), -1)
        loc, scale = self.x_nn(z_concat)
        return dist.Normal(loc, scale).to_event(1)

    # 返回二进制特征的条件分布
    def x_dist_binary(self, z, zt, zy):
        z_concat = torch.cat((z,zt,zy), -1)
        logits = self.x2_nn(z_concat)
        return dist.Bernoulli(logits=logits).to_event(1)

    # 根据t选择y0或y1网络返回y的条件分布
    def y_dist(self, t, z, zy):
        # Parameters are not shared among t values.
        z_concat = torch.cat((z, zy), -1)
        params0 = self.y0_nn(z_concat)
        params1 = self.y1_nn(z_concat)
        t = t.bool()
        params = [torch.where(t, p1, p0) for p0, p1 in zip(params0, params1)]
        return self.y0_nn.make_dist(*params)

    # # 返回t的条件分布
    def t_dist(self, z,zt):
        z_concat = torch.cat((z,zt), -1)
        logits, = self.t_nn(z_concat)
        return dist.Bernoulli(logits=logits)

class TEDVAE(nn.Module):
    def __init__(self, feature_dim, continuous_dim, binary_dim, outcome_dist="normal", 
                 latent_dim=20, latent_dim_t=20, latent_dim_y=20 , hidden_dim=200, num_layers=3, num_samples=100):
        # 将参数组织成字典
        config = dict(feature_dim=feature_dim, latent_dim=latent_dim,
                      latent_dim_t = latent_dim_t, latent_dim_y = latent_dim_y,
                      hidden_dim=hidden_dim, num_layers=num_layers, continuous_dim = continuous_dim, binary_dim = binary_dim,
                      num_samples=num_samples)
        # for name, size in config.items():
        #     if not (isinstance(size, int) and size > 0):
        #         raise ValueError("Expected {} > 0 but got {}".format(name, size))
        config["outcome_dist"] = outcome_dist
        self.feature_dim = feature_dim
        self.num_samples = num_samples

        super().__init__()
        self.model = Model(config)
        # Guide类,用于变分推断
        self.guide = Guide(config)
        # self.to_dev
        self.cuda()
    def fit(self, x, t, y,
            num_epochs=100,
            batch_size=100,
            learning_rate=1e-3,
            learning_rate_decay=0.1,
            weight_decay=1e-4):
        """
        Train using :class:`~pyro.infer.svi.SVI` with the
        :class:`TraceCausalEffect_ELBO` loss.

        :param ~torch.Tensor x:
        :param ~torch.Tensor t:
        :param ~torch.Tensor y:
        :param int num_epochs: Number of training epochs. Defaults to 100.
        :param int batch_size: Batch size. Defaults to 100.
        :param float learning_rate: Learning rate. Defaults to 1e-3.
        :param float learning_rate_decay: Learning rate decay over all epochs;
            the per-step decay rate will depend on batch size and number of epochs
            such that the initial learning rate will be ``learning_rate`` and the final
            learning rate will be ``learning_rate * learning_rate_decay``.
            Defaults to 0.1.
        :param float weight_decay: Weight decay. Defaults to 1e-4.
        :return: list of epoch losses
        """
        assert x.dim() == 2 and x.size(-1) == self.feature_dim
        assert t.shape == x.shape[:1]
        assert y.shape == y.shape[:1]
        # self.whiten = PreWhitener(x)

        # 创建数据集和数据加载器
        dataset = TensorDataset(x, t, y)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
        logger.info("Training with {} minibatches per epoch".format(len(dataloader)))
        num_steps = num_epochs * len(dataloader)

        # 使用 SVI（Stochastic Variational Inference）优化 ELBO 损失函数。
        # 使用 ClippedAdam 优化器和自定义的 ELBO（TraceCausalEffect_ELBO），它加入了额外的因果项。
        optim = ClippedAdam({"lr": learning_rate,
                             "weight_decay": weight_decay,
                             "lrd": learning_rate_decay ** (1 / num_steps)})
        svi = SVI(self.model, self.guide, optim, TraceCausalEffect_ELBO())
        losses = []
        # 训练循环，记录损失
        for epoch in range(num_epochs):
            for x, t, y in dataloader:
                # x = self.whiten(x)
                loss = svi.step(x, t, y, size=len(dataset)) / len(dataset)
                # print(loss)
                logger.debug("step {: >5d} loss = {:0.6g}".format(len(losses), loss))
                assert not torch_isnan(loss)
                losses.append(loss)
        return losses


    # 计算个体治疗效应(ITE)
    # 使用do算子计算反事实结果
    # 对t=0和t=1分别计算y的期望
    # ITE定义为y1 - y0的均值
    @torch.no_grad()
    def ite(self, x, ym, ys, num_samples=None, batch_size=None):
        r"""
        Computes Individual Treatment Effect for a batch of data ``x``.

        .. math::

            ITE(x) = \mathbb E\bigl[ \mathbf y \mid \mathbf X=x, do(\mathbf t=1) \bigr]
                   - \mathbb E\bigl[ \mathbf y \mid \mathbf X=x, do(\mathbf t=0) \bigr]

        This has complexity ``O(len(x) * num_samples ** 2)``.

        :param ~torch.Tensor x: A batch of data.
        :param int num_samples: The number of monte carlo samples.
            Defaults to ``self.num_samples`` which defaults to ``100``.
        :param int batch_size: Batch size. Defaults to ``len(x)``.
        :return: A ``len(x)``-sized tensor of estimated effects.
        :rtype: ~torch.Tensor
        """
        if num_samples is None:
            num_samples = self.num_samples
        if not torch._C._get_tracing_state():
            assert x.dim() == 2 and x.size(-1) == self.feature_dim

        dataloader = [x] if batch_size is None else DataLoader(x, batch_size=batch_size)
        logger.info("Evaluating {} minibatches".format(len(dataloader)))
        result = []
        for x in dataloader:
            # x = self.whiten(x)
            with pyro.plate("num_particles", num_samples, dim=-2):
                # 在 guide 中采样潜变量，记录其 trace。
                with poutine.trace() as tr, poutine.block(hide=["y", "t"]):
                    self.guide(x)
                with poutine.do(data=dict(t=torch.zeros(()))):
                    y0 = poutine.replay(self.model.y_mean, tr.trace)(x) * ys + ym
                with poutine.do(data=dict(t=torch.ones(()))):
                    y1 = poutine.replay(self.model.y_mean, tr.trace)(x) * ys + ym
            ite = (y1 - y0).mean(0)
            if not torch._C._get_tracing_state():
                logger.debug("batch ate = {:0.6g}".format(ite.mean()))
            result.append(ite)
        return torch.cat(result)

    # 用 torch.jit.trace_module 编译 ite 函数为 ScriptModule，便于部署。
    def to_script_module(self):
        """
        Compile this module using :func:`torch.jit.trace_module` ,
        assuming self has already been fit to data.

        :return: A traced version of self with an :meth:`ite` method.
        :rtype: torch.jit.ScriptModule
        """
        self.train(False)
        fake_x = torch.randn(2, self.feature_dim)
        with pyro.validation_enabled(False):
            # Disable check_trace due to nondeterministic nodes.
            result = torch.jit.trace_module(self, {"ite": (fake_x,)}, check_trace=False)
        return result
